//=========================================================================
// Copyright (C) 2024 The C++ Component Model(COMO) Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//=========================================================================

/**
 * Reliable UDP
 * https://github.com/cloudwu/rudp
 */

#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "rudp.h"

/**
 * RudpNew
 */
Rudp *RudpNew(Rudp *rudp, int send_delay, int expired_time)
{
    if (NULL == rudp) {
        rudp = malloc(sizeof(Rudp));
        if (NULL == rudp) {
            return NULL;
        }
    }

    memset(rudp, 0, sizeof(Rudp));

#ifdef COMO_FUNCTION_SAFETY_RTOS
    rudp->send_again.cap = MAX_SEND_AGAIN_ARRAY_SIZE;

    /*
    // The work below is unnecessary because memset() has set them to 0:
    for(int i = 0;  i < MAX_PACKAGE_BUFFER_SIZE;  i++) {
        rudp->usedBuffer[i] = false;
    }
    // Also, all int members has been initialized to 0.
    */
#endif

    rudp->send_delay = send_delay;
    rudp->expired = expired_time;

    return rudp;
}

#ifdef COMO_FUNCTION_SAFETY_RTOS
static int Pointer_in_packageBuffer(Rudp *U, void *m)
{
    if ((m < (void *)&(U->packageBuffer[0])) ||
                  (m >= (void *)&(U->packageBuffer[MAX_PACKAGE_BUFFER_SIZE]))) {
        return -1;
    }

    return (m - (void *)&(U->packageBuffer[0])) / MAX_PACKAGE_LENGTH;
}
#endif

/**
 * Clears the pointer used to return the packet
 */
static void clear_outpackage(Rudp *U)
{
    RudpPackage *tmp = U->send_package;
    while (NULL != tmp) {
        RudpPackage *next = tmp->next;
#ifdef COMO_FUNCTION_SAFETY_RTOS
        int i = Pointer_in_packageBuffer(U, tmp);
        if (i >= 0) {
            U->usedBuffer[i] = false;
        }
        else {
            // BUG
            assert(true);
        }
#else
        free(tmp);
#endif
        tmp = next;
    }
    U->send_package = NULL;
}

static void free_message_list(Rudp *U, struct rudp_message *m)
{
    while (NULL != m) {
        struct rudp_message *next = m->next;
#ifdef COMO_FUNCTION_SAFETY_RTOS
        int i = Pointer_in_packageBuffer(U, m);
        if (i >= 0) {
            U->usedBuffer[i] = false;
        }
        else {
            // BUG
            assert(true);
        }
        /*
        for(int i = 0;  i < MAX_PACKAGE_BUFFER_SIZE;  i++) {
            if (m == (struct rudp_message *)&(U->packageBuffer[i])) {
                U->usedBuffer[i] = false;
                break;
            }
        }*/
#else
        free(m);
#endif
        m = next;
    }
}

void RudpDelete(Rudp *U)
{
    free_message_list(U, U->send_queue.head);
    free_message_list(U, U->recv_queue.head);
    free_message_list(U, U->send_history.head);
    free_message_list(U, U->free_list);
    clear_outpackage(U);

#ifndef COMO_FUNCTION_SAFETY_RTOS
    free(U->send_again.a);

    // U is not necessarily from malloc()
    // free(U);
#endif
}

static struct rudp_message *new_message(Rudp *U, const uint8_t *buffer, int sz)
{
    // Attempt to get a free node from the free_list
    struct rudp_message *msg = U->free_list;

    // 1. If there is a free node
    if (NULL != msg) {
        U->free_list = msg->next;
        // Check if the node's buffer capacity meets the requirements
        if (msg->cap < sz) {
#ifdef COMO_FUNCTION_SAFETY_RTOS
            /**
             * This cannot happen when in function safety computing
             * COMO_FUNCTION_SAFETY_RTOS because the buffer is not dynamically
             * applied memory.
             */
            return NULL;
#else
            free(msg);
#endif
            // NOTICE: set to NULL to handle it later in case 2 
            msg = NULL;
        }
    }
    // 2. If no free node, or node's buffer capacity is insufficient
    if (NULL == msg) {
        int cap = sz;
        // Ensure the buffer capacity enough
        if (cap < GENERAL_PACKAGE) {
            cap = GENERAL_PACKAGE;
        }
// For FSCP, use a pre-allocated buffer to avoid dynamic memory allocation
#ifdef COMO_FUNCTION_SAFETY_RTOS
        msg = NULL;
        // Find an unused buffer
        for (int i = 0;  i < MAX_PACKAGE_BUFFER_SIZE;  i++) {
            if (! U->usedBuffer[i]) {
                msg = (struct rudp_message *)&(U->packageBuffer[i]);
                U->usedBuffer[i] = true;
                break;
            }
        }
#else
        // Dynamically allocate memory for a new node
        tmp = malloc(sizeof(struct rudp_message) + cap);
#endif
        // If memory allocation fails, return NULL
        if (NULL == msg) {
            return NULL;
        }
        msg->cap = cap;
    }
    // Init the node
    msg->sz = sz;
    msg->buffer = (uint8_t *)(msg + 1);
    // If there is data, copy it into the node's buffer
    if ((sz > 0) && (NULL != buffer)) {
        memcpy(msg->buffer, buffer, sz);
    }
    msg->tick = 0;
    msg->id = 0;
    msg->next = NULL;

    // Return a pointer to the new node
    return msg;
}

static void delete_message(Rudp *U, struct rudp_message *m)
{
    m->next = U->free_list;
    U->free_list = m;
}

static void queue_push(struct rudp_message_queue *q, struct rudp_message *m)
{
    if (NULL == q->tail) {
        q->head = m;
        q->tail = m;
    }
    else {
        q->tail->next = m;
        q->tail = m;
    }
}

static struct rudp_message *queue_pop(struct rudp_message_queue *q, int id)
{
    if (NULL == q->head) {
        return NULL;
    }

    struct rudp_message *m = q->head;
    if (m->id != id) {
        return NULL;
    }

    q->head = m->next;
    m->next = NULL;
    if (NULL == q->head) {
        q->tail = NULL;
    }
    return m;
}

static int array_insert(struct rudp_array *a, int id)
{
    int i;
    for (i = 0;  i < a->n;  i++) {
        if (a->a[i] == id) {
            return 0;
        }
        if (a->a[i] > id) {
            break;
        }
    }

#ifdef COMO_FUNCTION_SAFETY_RTOS
    // Memory pool not big enough
    if (a->n >= a->cap) {
        return -2;
    }
#else
    // insert before i
    if (a->n >= a->cap) {
        if (0 == a->cap) {
            a->cap = 16;
        }
        else {
            a->cap *= 2;
        }
        a->a = realloc(a->a, sizeof(int) * a->cap);
        if (NULL == a->a) {
            return -1;
        }
    }
#endif

    for (int j = a->n;  j > i;  j--) {
        a->a[j] = a->a[j - 1];
    }
    a->a[i] = id;
    a->n++;

    return 0;
}

/**
 * Send data in buffer
 */
int RudpSend(Rudp *U, const char *buffer, int sz)
{
    // Parameter validation 
    if ((sz <= 0) || (sz >= MAX_PACKAGE) || (NULL == buffer) || (NULL == U)) {
        return -1;
    }

    // Create an RUDP message
    struct rudp_message *m = new_message(U, (const uint8_t *)buffer, sz);
    if (NULL == m) {
        return -2;
    }

    // Set the message ID & timestamp
    m->id = U->send_id;
    m->tick = U->current_tick;
    // Update the send ID
    U->send_id++;

    // Push the message into the send queue
    queue_push(&(U->send_queue), m);

    return 0;
}

/**
 * Receive data into buffer
 */
int RudpRecv(Rudp *U, char *buffer)
{
    if ((NULL == buffer) || (NULL == U)) {
        return -1;
    }

    if (U->corrupt) {
        U->corrupt = 0;
        return -2;
    }
    struct rudp_message *tmp = queue_pop(&(U->recv_queue), U->recv_id_min);
    if (NULL == tmp) {
        return 0;
    }

    int sz = tmp->sz;

    U->recv_id_min++;
    if ((sz > 0) && (sz < MAX_PACKAGE)) {
        memcpy(buffer, tmp->buffer, sz);
    }
    else {
        return -3;
    }
    delete_message(U, tmp);

    return sz;
}

static void clear_send_expired(Rudp *U, int tick)
{
    struct rudp_message *m = U->send_history.head;
    struct rudp_message *last = NULL;

    /**
     * Traverse the list until we find a message with a tick greater than or
     * equal to the given tick.
     */
    while (NULL != m) {
        if (m->tick >= tick) {
            break;
        }
        last = m;
        m = m->next;
    }
    
    if (NULL != last) {
        // free all the messages before tick
        last->next = U->free_list;
        U->free_list = U->send_history.head;
    }
    U->send_history.head = m;
    if (NULL == m) {
        U->send_history.tail = NULL;
    }
}

static int get_id(Rudp *U, const uint8_t *buffer)
{
    int id = (int)buffer[0];
    id = id * 256 + buffer[1];

    id |= (U->recv_id_max & ~0xffffu);

    /**
     * The following constants (0x8000, 0x10000) are signed numbers,
     * if changed to signed numbers, will cause arithmetic errors.
     */

    if (id < (U->recv_id_max - 0x8000)) {
        id += 0x10000;
    }
    else if (id > (U->recv_id_max + 0x8000)) {
        id -= 0x10000;
    }
    return id;
}

static int add_request(Rudp *U, int id)
{
    return array_insert(&U->send_again, id);
}

static int insert_message(Rudp *U, int id, const uint8_t *buffer, int sz)
{
    if (id < U->recv_id_min) {
        return -1;
    }

    if ((id > U->recv_id_max) || (NULL == U->recv_queue.head)) {
        struct rudp_message *m = new_message(U, buffer, sz);
        if (NULL == m) {
            return -2;
        }

        m->id = id;
        queue_push(&U->recv_queue, m);
        U->recv_id_max = id;
    }
    else {
        struct rudp_message *m = U->recv_queue.head;
        struct rudp_message **last = &U->recv_queue.head;
        do {
            if (m->id == id) {
                return 0;
            }

            if (m->id > id) {
                // insert here
                struct rudp_message *tmp = new_message(U, buffer, sz);
                if (NULL == tmp) {
                    return -3;
                }

                tmp->id = id;
                tmp->next = m;
                *last = tmp;
                return 0;
            }
            last = &m->next;
            m = m->next;
        } while (m);
    }

    return 0;
}

static int add_missing(Rudp *U, int id)
{
    return insert_message(U, id, NULL, -1);
}

#define TYPE_IGNORE     0
#define TYPE_CORRUPT    1
#define TYPE_REQUEST    2
#define TYPE_MISSING    3
#define TYPE_NORMAL     4

/**
 * Extract package
 */
static int extract_package(Rudp *U, const uint8_t *buffer, int sz)
{
    int iRet;

    while (sz > 0) {
        int len = (int)buffer[0];
        if (len > 127) {
            if (sz <= 1) {
                U->corrupt = 1;
                return 0;
            }
            len = (len * 256 + buffer[1]) & 0x7fffu;
            buffer += 2;
            sz -= 2;
        }
        else {
            buffer += 1;
            sz -= 1;
        }

        switch (len) {
            case TYPE_IGNORE:
                if (0 == U->send_again.n) {
                    // request next package id
                    iRet = array_insert(&U->send_again, U->recv_id_min);
                    if (iRet < 0) {
                        return -1;
                    }
                }
                break;

            case TYPE_CORRUPT:
                U->corrupt = 1;
                return 0;

            case TYPE_REQUEST:
            case TYPE_MISSING:
                if (sz < 2) {
                    U->corrupt = 1;
                    return 0;
                }

                if (TYPE_REQUEST == len) {
                    iRet = add_request(U, get_id(U, buffer));
                    if (iRet < 0) {
                        return -2;
                    }
                }
                else {
                    iRet = add_missing(U, get_id(U, buffer));
                    if (iRet < 0) {
                        return -3;
                    }
                }

                buffer += 2;
                sz -= 2;
                break;

            default:
                len -= TYPE_NORMAL;
                if (sz < (len + 2)) {
                    U->corrupt = 1;
                    return 0;
                }
                else {
                    int id = get_id(U, buffer);
                    iRet = insert_message(U, id, buffer + 2, len);
                    if (iRet < 0) {
                        return -4;
                    }
                }
                buffer += len + 2;
                sz -= len + 2;
                break;
        }
    }

    return 0;
}

struct tmp_buffer {
    uint8_t buf[GENERAL_PACKAGE];
    int sz;
    RudpPackage *head;
    RudpPackage *tail;
};

static bool new_package(Rudp *U, struct tmp_buffer *tmp)
{
#ifdef COMO_FUNCTION_SAFETY_RTOS
    RudpPackage *p = NULL;
    for(int i = 0;  i < MAX_PACKAGE_BUFFER_SIZE;  i++) {
        if (! U->usedBuffer[i]) {
            p = (RudpPackage *)&(U->packageBuffer[i]);
            U->usedBuffer[i] = true;
            break;
        }
    }
#else
    RudpPackage *p = malloc(sizeof(RudpPackage) + tmp->sz);
#endif

    if (NULL == p) {
        return false;
    }

    p->next = NULL;
    p->buffer = (char *)(p + 1);
    p->sz = tmp->sz;
    memcpy(p->buffer, tmp->buf, tmp->sz);
    if (NULL == tmp->tail) {
        tmp->head = p;
        tmp->tail = p;
    }
    else {
        tmp->tail->next = p;
        tmp->tail = p;
    }
    tmp->sz = 0;

    return true;
}

static int fill_header(uint8_t *buf, int len, int id)
{
    int sz;

    if (len < 128) {
        buf[0] = len;
        ++buf;
        sz = 1;
    }
    else {
        buf[0] = ((len & 0x7f00u) >> 8) | 0x80u;
        buf[1] = len & 0xffu;
        buf += 2;
        sz = 2;
    }
    buf[0] = (id & 0xff00u) >> 8;
    buf[1] = id & 0xffu;

    return sz + 2;
}

static bool pack_request(Rudp *U, struct tmp_buffer *tmp, int id, int tag)
{
    int sz = GENERAL_PACKAGE - tmp->sz;
    if (sz < 3) {
        if (! new_package(U, tmp)) {
            return false;
        }
    }
    uint8_t *buffer = tmp->buf + tmp->sz;
    tmp->sz += fill_header(buffer, tag, id);

    return true;
}

/**
 * Pack message into temporary buffer
 * 
 * Checks if the message size exceeds the remaining space in the temporary buffer,
 * if so, try to create a new package. Handle too large messages specially.
 * 
 * @param U     RUDP instance pointer
 * @param tmp   Temporary buffer
 * @param m     Message pointer, containing the message ID, buffer, size...
 * @return      True on success, false on failure
 */
static bool pack_message(Rudp *U, struct tmp_buffer *tmp, struct rudp_message *m)
{
    int sz = GENERAL_PACKAGE - tmp->sz;
    // Check if the message size > the remaining space in the temporary buffer
    if (m->sz > (GENERAL_PACKAGE - 4)) {
        // If there are already messages in the tmp buffer, create a new package
        if (tmp->sz > 0) {
            if (! new_package(U, tmp)) {
                return false;
            }
        }

        // Big package
        sz = 4 + m->sz;
#ifdef COMO_FUNCTION_SAFETY_RTOS
        // Under RTOS environment, find an unused package buffer
        RudpPackage *p = NULL;
        for(int i = 0;  i < MAX_PACKAGE_BUFFER_SIZE;  i++) {
            if (! U->usedBuffer[i]) {
                p = (RudpPackage *)&(U->packageBuffer[i]);
                U->usedBuffer[i] = true;
                break;
            }
        }
#else
        // In a non-RTOS environment, dynamically allocate package memory
        RudpPackage *p = malloc(sizeof(RudpPackage) + sz);
#endif
        // Check if memory allocation failed
        if (NULL == p) {
            return false;
        }

        // Init package
        p->next = NULL;
        p->buffer = (char *)(p + 1);
        p->sz = sz;

        // Fill the package header with message size and ID
        fill_header((uint8_t *)p->buffer, m->sz + TYPE_NORMAL, m->id);
        // Copy message content to package
        memcpy(p->buffer + 4, m->buffer, m->sz);
        // If tmp buffer is empty, set the head and tail to the new package
        if (NULL == tmp->tail) {
            tmp->head = p;
            tmp->tail = p;
        }
        // Otherwise, link the new package to the tail
        else {
            tmp->tail->next = p;
            tmp->tail = p;
        }
        return true;
    }
    // If remaining space in tmp buffer < message size, create a new package
    if (sz < (4 + m->sz)) {
        if (! new_package(U, tmp)) {
            return false;
        }
    }
    // Calculate the position to store the new message in the tmp buffer
    uint8_t *buf = tmp->buf + tmp->sz;
    // Fill the message header and update the tmp buffer size
    int len = fill_header(buf, m->sz + TYPE_NORMAL, m->id);
    tmp->sz += len + m->sz;
    buf += len;
    // Copy message content to temporary buffer
    memcpy(buf, m->buffer, m->sz);

    return true;
}

/**
 * Request missing messages by checking message ID in the receive queue.
 * @param U     Rudp pointer
 * @param tmp   Temporary buffer to be edited
 * @return      True on success, false on failure
 */
static bool request_missing(Rudp *U, struct tmp_buffer *tmp)
{
    int id = U->recv_id_min;
    struct rudp_message *m = U->recv_queue.head;
    while (NULL != m) {
        assert(m->id >= id);

        if (m->id > id) {
            for (int i = id;  i < m->id;  i++) {
                if (! pack_request(U, tmp, i, TYPE_REQUEST)) {
                    return false;
                }
            }
        }
        id = m->id + 1;
        m = m->next;
    }

    return true;
}

static bool reply_request(Rudp *U, struct tmp_buffer *tmp)
{
    struct rudp_message *history = U->send_history.head;
    for (int i = 0;  i < U->send_again.n;  i++) {
        int id = U->send_again.a[i];
        if (id < U->recv_id_min) {
            // alreay recv, ignore
            continue;
        }
        for (;;) {
            if ((NULL == history) || (id < history->id)) {
                // expired
                if (! pack_request(U, tmp, id, TYPE_MISSING)) {
                    return false;
                }
                break;
            }
            else if (id == history->id) {
                if (! pack_message(U, tmp, history)) {
                    return false;
                }
                break;
            }
            history = history->next;
        }
    }

    U->send_again.n = 0;

    return true;
}

static bool send_message(Rudp *U, struct tmp_buffer *tmp)
{
    struct rudp_message *m = U->send_queue.head;
    while (NULL != m) {
        // Pack each message into the temporary buffer
        if (! pack_message(U, tmp, m)) {
            return false;
        }
        m = m->next;
    }
    if (NULL != U->send_queue.head) {
        if (NULL == U->send_history.tail) {
            U->send_history = U->send_queue;
        }
        else {
            U->send_history.tail->next = U->send_queue.head;
            U->send_history.tail = U->send_queue.tail;
        }
        U->send_queue.head = NULL;
        U->send_queue.tail = NULL;
    }

    return true;
}


/**
 * 1. request missing ( lookup U->recv_queue )
 * 2. reply request ( U->send_again )
 * 3. send message ( U->send_queue )
 * 4. send heartbeat
 */
static RudpPackage *gen_outpackage(Rudp *U)
{
    struct tmp_buffer tmp;
    tmp.sz = 0;
    tmp.head = NULL;
    tmp.tail = NULL;

    bool tf;
    tf  = request_missing(U, &tmp);
    tf &= reply_request(U, &tmp);
    tf &= send_message(U, &tmp);
    if (! tf) {
        return NULL;
    }

    // Close tmp
    if (NULL == tmp.head) {
        if (0 == tmp.sz) {
            tmp.buf[0] = TYPE_IGNORE;
            tmp.sz = 1;
        }
    }
    if (! new_package(U, &tmp)) {
        return NULL;
    }

    return tmp.head;
}

/**
 * RudpUpdate: Updates the Rudp connection status and processes data packets
 * @param U       Pointer to the Rudp object, representing the connection context
 * @param buffer  Pointer to the data packet buffer
 * @param sz      Size of the data packet
 * @param tick    Time interval since last call, to update the connection status
 * @return        Pointer to the RudpPackage to be sent
 */
RudpPackage *RudpUpdate(Rudp *U, const void *buffer, int sz, int tick)
{
    int iRet;

    // Update the current time point of the connection
    U->current_tick += tick;

    // Clear packages
    clear_outpackage(U);

    // Extract and process the received data packet
    iRet = extract_package(U, buffer, sz);
    if (iRet < 0) {
        return NULL;
    }

    // The time of sending and receiving the last data packet has passed

    // Check if the time since the last expiration has exceeded the configured expiration duration
    if (U->current_tick >= (U->last_expired_tick + U->expired)) {
        clear_send_expired(U, U->last_expired_tick);
        U->last_expired_tick = U->current_tick;
    }

    // Check if it's time to send a new package based on the send delay
    if (U->current_tick >= (U->last_send_tick + U->send_delay)) {
        U->send_package = gen_outpackage(U);
        U->last_send_tick = U->current_tick;

        return U->send_package;
    }

    return NULL;
}
